import { loadGraphModel } from "@tensorflow/tfjs-converter";
import { ready, io } from "@tensorflow/tfjs-core";
import { FaceLandmarksDetection } from "./face";

export { FaceLandmarksDetection };
export type { NormalizedFace, FacePrediction } from "./face";
export interface OptionTypes {
  modelUrl: string;
  maxFaces?: number;
  inputWidth?: number;
  inputHeight?: number;
  iouThreshold?: number;
  scoreThreshold?: number;
  loadOptions?: io.LoadOptions;
}
export const load = async ({
  maxFaces = 10,
  inputWidth = 128,
  inputHeight = 128,
  iouThreshold = 0.3,
  scoreThreshold = 0.75,
  loadOptions = {
    fromTFHub: true,
  },
  modelUrl,
}: OptionTypes) => {
  // tfjs ready
  await ready();
  // load graph model
  const faceGraphModel = await loadGraphModel(modelUrl, loadOptions);
  // face model
  const model = new FaceLandmarksDetection(
    faceGraphModel,
    inputWidth,
    inputHeight,
    maxFaces,
    iouThreshold,
    scoreThreshold,
  );
  return model;
};
